import pickle
import onnxruntime as ort
from onnx_tf.backend import prepare
import torch
from torch import nn


class Net(nn.Module):
    def __init__(self) -> None:
        super().__init__()

    def forward(self, x):
        out = torch.abs(input=x)
        return out


def run_pytorch_and_tensorflow():
    model = Net()
    torch.save(model, "")

    tensor = torch.randn(14, 0, 9, 12, 6)

    # transfer *.pth model to *.onnx model
    input_list = (tensor)
    input_names = ["tensor"]
    input_dict = np.load("/home/ubuntu/Ascend/dataset/pytorch/seed/torch.nn.init.orthogonal_/20221228_222523/torch.nn.init.orthogonal__seeds.npz")
    kwargs = {"tensor": torch.as_tensor(input_dict["tensor_1499"])}
    model = torch.load("/home/ubuntu/Ascend/models/pytorch/torch.nn.init.orthogonal_/20221228_222523/torch.nn.init.orthogonal__1499.pth")
    if torch.cuda.is_available():
        model = model.cuda()
        kwargs["tensor"] = kwargs["tensor"].cuda()
    model.eval()
    res = model(**kwargs).cpu()
    result_dict = {"torch.nn.init.orthogonal__1499": res}
    np.savez("/home/ubuntu/Ascend/results/pt2cann/compare_results_gpu/torch.nn.init.orthogonal_/20221228_222523/torch.nn.init.orthogonal__1499_results", **result_dict)
